import numpy as np
from collections import Counter


# 朴素贝叶斯分类器类
class NaiveBayes:
    def __init__(self):
        # 存储每个类别的先验概率
        self.class_priors = {}
        # 存储每个特征在给定类别下的条件概率
        self.feature_likelihoods = {}
        # 存储所有可能的类别标签
        self.classes = []

    def fit(self, X, y):
        """
        训练模型：计算先验概率和条件概率。

        参数:
        X: 训练数据集，二维数组，形状为 (样本数量, 特征数量)
        y: 样本标签，一维数组
        """
        self.classes = np.unique(y)  # 获取所有类别
        num_samples, num_features = X.shape  # 获取样本数量和特征数量

        # 计算每个类别的先验概率 P(C)
        for c in self.classes:
            self.class_priors[c] = np.mean(y == c)

        # 对于每个类别c，计算每个特征的条件概率 P(X|C)
        for c in self.classes:
            X_c = X[y == c]  # 筛选出属于类别c的样本
            likelihoods = {}
            for feature_idx in range(num_features):
                # 统计每个特征值出现的次数
                feature_counts = Counter(X_c[:, feature_idx])
                # 计算条件概率，加上平滑处理避免零概率
                total_count = sum(feature_counts.values()) + len(feature_counts.keys())
                likelihoods[feature_idx] = {feat_val: (count + 1) / total_count
                                            for feat_val, count in feature_counts.items()}
            self.feature_likelihoods[c] = likelihoods

    def predict(self, X):
        """
        预测新样本的类别标签。

        参数:
        X: 测试数据集，二维数组

        返回:
        predictions: 预测结果，一维数组
        """
        predictions = []
        for sample in X:
            posteriors = {}
            for c in self.classes:
                # 初始后验概率为先验概率的对数（防止浮点下溢）
                posterior = np.log(self.class_priors[c])
                for feature_idx, feature_value in enumerate(sample):
                    # 如果特征值未在训练集中出现过，则使用一个很小的默认概率
                    likelihood = self.feature_likelihoods[c].get(feature_idx, {}).get(feature_value, 1e-6)
                    posterior += np.log(likelihood)  # 增加条件概率的对数值
                posteriors[c] = posterior
            # 选择具有最高后验概率的类别作为预测结果
            predictions.append(max(posteriors, key=posteriors.get))
        return np.array(predictions)


# 数据加载与预处理函数（仅用于示例，实际使用时应根据具体情况调整）
def load_data():
    # 这里为了简化，仍然使用了sklearn的数据集加载功能，但在实际项目中应该用其他方法获取数据。
    from sklearn.datasets import fetch_20newsgroups
    from sklearn.feature_extraction.text import CountVectorizer
    from sklearn.model_selection import train_test_split

    # 加载20Newsgroups数据集中的两个类别
    newsgroups = fetch_20newsgroups(subset='all', categories=['sci.space', 'rec.autos'])
    vectorizer = CountVectorizer(binary=True, stop_words='english', max_features=1000)
    X = vectorizer.fit_transform(newsgroups.data).toarray()
    y = newsgroups.target

    # 分割数据集为训练集和测试集
    return train_test_split(X, y, test_size=0.2, random_state=42)


# 主函数入口
if __name__ == "__main__":
    X_train, X_test, y_train, y_test = load_data()

    # 创建并训练朴素贝叶斯分类器
    model = NaiveBayes()
    model.fit(X_train, y_train)

    # 使用训练好的模型进行预测
    y_pred = model.predict(X_test)

    # 输出准确率
    accuracy = np.mean(y_pred == y_test)
    print(f"测试集准确率: {accuracy:.4f}")